{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mq-riZs-TJGt"
      },
      "source": [
        "##### Copyright 2024 The AI Edge Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LEvnopDoTC4M"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QSRG6qmtTRSk"
      },
      "source": [
        "# TensorFlow Lite Metadata Writer API\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JlzjEt4Txr0x"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/lite/models/convert/metadata_writer_tutorial\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/convert/metadata_writer_tutorial.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/g3doc/models/convert/metadata_writer_tutorial.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/tensorflow/tensorflow/lite/g3doc/models/convert/metadata_writer_tutorial.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b0gwEhfRYat6"
      },
      "source": [
        "[TensorFlow Lite Model Metadata](https://www.tensorflow.org/lite/models/convert/metadata) is a standard model description format. It contains rich semantics for general model information, inputs/outputs, and associated files, which makes the model self-descriptive and exchangeable.\n",
        "\n",
        "Model Metadata is currently used in the following two primary use cases:\n",
        "1. **Enable easy model inference using TensorFlow Lite [Task Library](https://www.tensorflow.org/lite/inference_with_metadata/task_library/overview) and [codegen tools](https://www.tensorflow.org/lite/inference_with_metadata/codegen)**. Model Metadata contains the mandatory information required during inference, such as label files in image classification, sampling rate of the audio input in audio classification, and tokenizer type to process input string in Natural Language models.\n",
        "\n",
        "2. **Enable model creators to include documentation**, such as the description of model inputs/outputs or how to use the model. Model users can view this documentation via visualization tools such as [Netron](https://netron.app/).\n",
        "\n",
        "TensorFlow Lite Metadata Writer API provides an easy-to-use API to create Model Metadata for popular ML tasks supported by the TFLite Task Library.  This notebook shows examples of how the metadata should be populated for the following tasks below:\n",
        "\n",
        "*   [Image classifiers](#image_classifiers)\n",
        "*   [Object detectors](#object_detectors)\n",
        "*   [Image segmenters](#image_segmenters)\n",
        "*   [Natural language classifiers](#nl_classifiers)\n",
        "*   [Audio classifiers](#audio_classifiers)\n",
        "\n",
        "Metadata writers for BERT natural language classifiers and BERT question answerers are coming soon.\n",
        "\n",
        "If you want to add metadata for use cases that are not supported, please use the [Flatbuffers Python API](https://www.tensorflow.org/lite/models/convert/metadata#adding_metadata). See the tutorials [here](https://www.tensorflow.org/lite/models/convert/metadata#adding_metadata).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GVRIGdA4T6tO"
      },
      "source": [
        "## Prerequisites"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bVTD2KSyotBK"
      },
      "source": [
        "Install the TensorFlow Lite Support Pypi package."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m-8xSrSvUg-6"
      },
      "outputs": [],
      "source": [
        "!pip install tflite-support-nightly"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hyYS87Odpxef"
      },
      "source": [
        "## Create Model Metadata for Task Library and Codegen"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uLxv541TqTim"
      },
      "source": [
        "\u003ca name=image_classifiers\u003e\u003c/a\u003e\n",
        "### Image classifiers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s41TjCGlsyEF"
      },
      "source": [
        "See the [image classifier model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_classifier#model_compatibility_requirements) for more details about the supported model format."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_KsPKmg8T9-8"
      },
      "source": [
        "Step 1: Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hhgNqEtWrwB3"
      },
      "outputs": [],
      "source": [
        "from tflite_support.metadata_writers import image_classifier\n",
        "from tflite_support.metadata_writers import writer_utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o9WBgiFdsiIQ"
      },
      "source": [
        "Step 2: Download the example image classifier, [mobilenet_v2_1.0_224.tflite](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite), and the [label file](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6WgSBbNet-Tt"
      },
      "outputs": [],
      "source": [
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite -o mobilenet_v2_1.0_224.tflite\n",
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt -o mobilenet_labels.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ALtlz7woweHe"
      },
      "source": [
        "Step 3: Create metadata writer and populate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_SMEBBt2r-W6"
      },
      "outputs": [],
      "source": [
        "ImageClassifierWriter = image_classifier.MetadataWriter\n",
        "_MODEL_PATH = \"mobilenet_v2_1.0_224.tflite\"\n",
        "# Task Library expects label files that are in the same format as the one below.\n",
        "_LABEL_FILE = \"mobilenet_labels.txt\"\n",
        "_SAVE_TO_PATH = \"mobilenet_v2_1.0_224_metadata.tflite\"\n",
        "# Normalization parameters are required when reprocessing the image. It is\n",
        "# optional if the image pixel values are in the range of [0, 255] and the input\n",
        "# tensor is quantized to uint8. See the introduction for normalization and\n",
        "# quantization parameters below for more details.\n",
        "# https://www.tensorflow.org/lite/models/convert/metadata#normalization_and_quantization_parameters)\n",
        "_INPUT_NORM_MEAN = 127.5\n",
        "_INPUT_NORM_STD = 127.5\n",
        "\n",
        "# Create the metadata writer.\n",
        "writer = ImageClassifierWriter.create_for_inference(\n",
        "    writer_utils.load_file(_MODEL_PATH), [_INPUT_NORM_MEAN], [_INPUT_NORM_STD],\n",
        "    [_LABEL_FILE])\n",
        "\n",
        "# Verify the metadata generated by the metadata writer.\n",
        "print(writer.get_metadata_json())\n",
        "\n",
        "# Populate the metadata into the model.\n",
        "writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GhhTDkr-uf0n"
      },
      "source": [
        "\u003ca name=object_detectors\u003e\u003c/a\u003e\n",
        "### Object detectors"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EL9GssnTuf0n"
      },
      "source": [
        "See the [object detector model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/object_detector#model_compatibility_requirements) for more details about the supported model format."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r-HUTEtHuf0n"
      },
      "source": [
        "Step 1: Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2_NIROeouf0o"
      },
      "outputs": [],
      "source": [
        "from tflite_support.metadata_writers import object_detector\n",
        "from tflite_support.metadata_writers import writer_utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UM6jijiUuf0o"
      },
      "source": [
        "Step 2: Download the example object detector, [ssd_mobilenet_v1.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/ssd_mobilenet_v1.tflite), and the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/labelmap.txt)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4i_BBfGzuf0o"
      },
      "outputs": [],
      "source": [
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/ssd_mobilenet_v1.tflite -o ssd_mobilenet_v1.tflite\n",
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/object_detector/labelmap.txt -o ssd_mobilenet_labels.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DG9T3eSDwsnd"
      },
      "source": [
        "Step 3: Create metadata writer and populate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vMGGeJfCuf0p"
      },
      "outputs": [],
      "source": [
        "ObjectDetectorWriter = object_detector.MetadataWriter\n",
        "_MODEL_PATH = \"ssd_mobilenet_v1.tflite\"\n",
        "# Task Library expects label files that are in the same format as the one below.\n",
        "_LABEL_FILE = \"ssd_mobilenet_labels.txt\"\n",
        "_SAVE_TO_PATH = \"ssd_mobilenet_v1_metadata.tflite\"\n",
        "# Normalization parameters are required when reprocessing the image. It is\n",
        "# optional if the image pixel values are in the range of [0, 255] and the input\n",
        "# tensor is quantized to uint8. See the introduction for normalization and\n",
        "# quantization parameters below for more details.\n",
        "# https://www.tensorflow.org/lite/models/convert/metadata#normalization_and_quantization_parameters)\n",
        "_INPUT_NORM_MEAN = 127.5\n",
        "_INPUT_NORM_STD = 127.5\n",
        "\n",
        "# Create the metadata writer.\n",
        "writer = ObjectDetectorWriter.create_for_inference(\n",
        "    writer_utils.load_file(_MODEL_PATH), [_INPUT_NORM_MEAN], [_INPUT_NORM_STD],\n",
        "    [_LABEL_FILE])\n",
        "\n",
        "# Verify the metadata generated by the metadata writer.\n",
        "print(writer.get_metadata_json())\n",
        "\n",
        "# Populate the metadata into the model.\n",
        "writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QT0Oa0SU6uGS"
      },
      "source": [
        "\u003ca name=image_segmenters\u003e\u003c/a\u003e\n",
        "### Image segmenters"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XaFQmg-S6uGW"
      },
      "source": [
        "See the [image segmenter model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/image_segmenter#model_compatibility_requirements) for more details about the supported model format."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DiktANhj6uGX"
      },
      "source": [
        "Step 1: Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H6Lrw3op6uGX"
      },
      "outputs": [],
      "source": [
        "from tflite_support.metadata_writers import image_segmenter\n",
        "from tflite_support.metadata_writers import writer_utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9EFs8Oyi6uGX"
      },
      "source": [
        "Step 2: Download the example image segmenter, [deeplabv3.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/deeplabv3.tflite), and the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/labelmap.txt)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "feQDH0bN6uGY"
      },
      "outputs": [],
      "source": [
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/deeplabv3.tflite -o deeplabv3.tflite\n",
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_segmenter/labelmap.txt -o deeplabv3_labels.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8LhiAbJM6uGY"
      },
      "source": [
        "Step 3: Create metadata writer and populate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yot8xLI46uGY"
      },
      "outputs": [],
      "source": [
        "ImageSegmenterWriter = image_segmenter.MetadataWriter\n",
        "_MODEL_PATH = \"deeplabv3.tflite\"\n",
        "# Task Library expects label files that are in the same format as the one below.\n",
        "_LABEL_FILE = \"deeplabv3_labels.txt\"\n",
        "_SAVE_TO_PATH = \"deeplabv3_metadata.tflite\"\n",
        "# Normalization parameters are required when reprocessing the image. It is\n",
        "# optional if the image pixel values are in the range of [0, 255] and the input\n",
        "# tensor is quantized to uint8. See the introduction for normalization and\n",
        "# quantization parameters below for more details.\n",
        "# https://www.tensorflow.org/lite/models/convert/metadata#normalization_and_quantization_parameters)\n",
        "_INPUT_NORM_MEAN = 127.5\n",
        "_INPUT_NORM_STD = 127.5\n",
        "\n",
        "# Create the metadata writer.\n",
        "writer = ImageSegmenterWriter.create_for_inference(\n",
        "    writer_utils.load_file(_MODEL_PATH), [_INPUT_NORM_MEAN], [_INPUT_NORM_STD],\n",
        "    [_LABEL_FILE])\n",
        "\n",
        "# Verify the metadata generated by metadata writer.\n",
        "print(writer.get_metadata_json())\n",
        "\n",
        "# Populate the metadata into the model.\n",
        "writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NnvM80e7AG-h"
      },
      "source": [
        "\u003ca name=nl_classifiers\u003e\u003c/a\u003e\n",
        "###Natural language classifiers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dfOPhFwOAG-k"
      },
      "source": [
        "See the [natural language classifier model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/nl_classifier#model_compatibility_requirements) for more details about the supported model format."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WMJ7tvuwAG-k"
      },
      "source": [
        "Step 1: Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_FGVyb2iAG-k"
      },
      "outputs": [],
      "source": [
        "from tflite_support.metadata_writers import nl_classifier\n",
        "from tflite_support.metadata_writers import metadata_info\n",
        "from tflite_support.metadata_writers import writer_utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iIg7rATpAG-l"
      },
      "source": [
        "Step 2: Download the example natural language classifier, [movie_review.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/movie_review.tflite), the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/labels.txt), and the [vocab file](https://storage.googleapis.com/download.tensorflow.org/models/tflite_support/nl_classifier/vocab.txt)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TzuQcti2AG-l"
      },
      "outputs": [],
      "source": [
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/movie_review.tflite -o movie_review.tflite\n",
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/nl_classifier/labels.txt -o movie_review_labels.txt\n",
        "!curl -L https://storage.googleapis.com/download.tensorflow.org/models/tflite_support/nl_classifier/vocab.txt -o movie_review_vocab.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BWxUtHdeAG-m"
      },
      "source": [
        "Step 3: Create metadata writer and populate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NGPWzRuHAG-m"
      },
      "outputs": [],
      "source": [
        "NLClassifierWriter = nl_classifier.MetadataWriter\n",
        "_MODEL_PATH = \"movie_review.tflite\"\n",
        "# Task Library expects label files and vocab files that are in the same formats\n",
        "# as the ones below.\n",
        "_LABEL_FILE = \"movie_review_labels.txt\"\n",
        "_VOCAB_FILE = \"movie_review_vocab.txt\"\n",
        "# NLClassifier supports tokenize input string using the regex tokenizer. See\n",
        "# more details about how to set up RegexTokenizer below:\n",
        "# https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/metadata_writers/metadata_info.py#L130\n",
        "_DELIM_REGEX_PATTERN = r\"[^\\w\\']+\"\n",
        "_SAVE_TO_PATH = \"moview_review_metadata.tflite\"\n",
        "\n",
        "# Create the metadata writer.\n",
        "writer = nl_classifier.MetadataWriter.create_for_inference(\n",
        "    writer_utils.load_file(_MODEL_PATH),\n",
        "    metadata_info.RegexTokenizerMd(_DELIM_REGEX_PATTERN, _VOCAB_FILE),\n",
        "    [_LABEL_FILE])\n",
        "\n",
        "# Verify the metadata generated by metadata writer.\n",
        "print(writer.get_metadata_json())\n",
        "\n",
        "# Populate the metadata into the model.\n",
        "writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qv0WDnzW711f"
      },
      "source": [
        "\u003ca name=audio_classifiers\u003e\u003c/a\u003e\n",
        "### Audio classifiers"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xqP7X8jww8pL"
      },
      "source": [
        "See the [audio classifier model compatibility requirements](https://www.tensorflow.org/lite/inference_with_metadata/task_library/audio_classifier#model_compatibility_requirements) for more details about the supported model format."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7RToKepxw8pL"
      },
      "source": [
        "Step 1: Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JjddvTXKw8pL"
      },
      "outputs": [],
      "source": [
        "from tflite_support.metadata_writers import audio_classifier\n",
        "from tflite_support.metadata_writers import metadata_info\n",
        "from tflite_support.metadata_writers import writer_utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ar418rH6w8pL"
      },
      "source": [
        "Step 2: Download the example audio classifier, [yamnet.tflite](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_wavin_quantized_mel_relu6.tflite), and the [label file](https://github.com/tensorflow/tflite-support/blob/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_521_labels.txt)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5eQY6znmw8pM"
      },
      "outputs": [],
      "source": [
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_wavin_quantized_mel_relu6.tflite -o yamnet.tflite\n",
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/audio_classifier/yamnet_521_labels.txt -o yamnet_labels.txt\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1TYP5w0Ew8pM"
      },
      "source": [
        "Step 3: Create metadata writer and populate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MDlSczBQw8pM"
      },
      "outputs": [],
      "source": [
        "AudioClassifierWriter = audio_classifier.MetadataWriter\n",
        "_MODEL_PATH = \"yamnet.tflite\"\n",
        "# Task Library expects label files that are in the same format as the one below.\n",
        "_LABEL_FILE = \"yamnet_labels.txt\"\n",
        "# Expected sampling rate of the input audio buffer.\n",
        "_SAMPLE_RATE = 16000\n",
        "# Expected number of channels of the input audio buffer. Note that Task library only\n",
        "# supports single channel so far.\n",
        "_CHANNELS = 1\n",
        "_SAVE_TO_PATH = \"yamnet_metadata.tflite\"\n",
        "\n",
        "# Create the metadata writer.\n",
        "writer = AudioClassifierWriter.create_for_inference(\n",
        "    writer_utils.load_file(_MODEL_PATH), _SAMPLE_RATE, _CHANNELS, [_LABEL_FILE])\n",
        "\n",
        "# Verify the metadata generated by the metadata writer.\n",
        "print(writer.get_metadata_json())\n",
        "\n",
        "# Populate the metadata into the model.\n",
        "writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YoRLs84yNAJR"
      },
      "source": [
        "## Create Model Metadata with semantic information"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cxXsOBknOGJ2"
      },
      "source": [
        "You can fill in more descriptive information about the model and each tensor through the Metadata Writer API to help improve model understanding. It can be done through the 'create_from_metadata_info' method in each metadata writer. In general, you can fill in data through the parameters of 'create_from_metadata_info', i.e. `general_md`, `input_md`, and `output_md`. See the example below to create a rich Model Metadata for image classifiers."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q-LW6nrcQ9lv"
      },
      "source": [
        "Step 1: Import the required packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KsL_egYcRGw3"
      },
      "outputs": [],
      "source": [
        "from tflite_support.metadata_writers import image_classifier\n",
        "from tflite_support.metadata_writers import metadata_info\n",
        "from tflite_support.metadata_writers import writer_utils\n",
        "from tflite_support import metadata_schema_py_generated as _metadata_fb"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0UWck_8uRboF"
      },
      "source": [
        "Step 2: Download the example image classifier, [mobilenet_v2_1.0_224.tflite](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite), and the [label file](https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TqJ-jh-PRVdk"
      },
      "outputs": [],
      "source": [
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/mobilenet_v2_1.0_224.tflite -o mobilenet_v2_1.0_224.tflite\n",
        "!curl -L https://github.com/tensorflow/tflite-support/raw/master/tensorflow_lite_support/metadata/python/tests/testdata/image_classifier/labels.txt -o mobilenet_labels.txt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r4I5wJMQRxzb"
      },
      "source": [
        "Step 3: Create model and tensor information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "urd7HDuaR_HC"
      },
      "outputs": [],
      "source": [
        "model_buffer = writer_utils.load_file(\"mobilenet_v2_1.0_224.tflite\")\n",
        "\n",
        "# Create general model information.\n",
        "general_md = metadata_info.GeneralMd(\n",
        "    name=\"ImageClassifier\",\n",
        "    version=\"v1\",\n",
        "    description=(\"Identify the most prominent object in the image from a \"\n",
        "                 \"known set of categories.\"),\n",
        "    author=\"TensorFlow Lite\",\n",
        "    licenses=\"Apache License. Version 2.0\")\n",
        "\n",
        "# Create input tensor information.\n",
        "input_md = metadata_info.InputImageTensorMd(\n",
        "    name=\"input image\",\n",
        "    description=(\"Input image to be classified. The expected image is \"\n",
        "                 \"128 x 128, with three channels (red, blue, and green) per \"\n",
        "                 \"pixel. Each element in the tensor is a value between min and \"\n",
        "                 \"max, where (per-channel) min is [0] and max is [255].\"),\n",
        "    norm_mean=[127.5],\n",
        "    norm_std=[127.5],\n",
        "    color_space_type=_metadata_fb.ColorSpaceType.RGB,\n",
        "    tensor_type=writer_utils.get_input_tensor_types(model_buffer)[0])\n",
        "\n",
        "# Create output tensor information.\n",
        "output_md = metadata_info.ClassificationTensorMd(\n",
        "    name=\"probability\",\n",
        "    description=\"Probabilities of the 1001 labels respectively.\",\n",
        "    label_files=[\n",
        "        metadata_info.LabelFileMd(file_path=\"mobilenet_labels.txt\",\n",
        "                                  locale=\"en\")\n",
        "    ],\n",
        "    tensor_type=writer_utils.get_output_tensor_types(model_buffer)[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N5aL5Uxkf4aO"
      },
      "source": [
        "Step 4: Create metadata writer and populate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_iWIwdqEf_mr"
      },
      "outputs": [],
      "source": [
        "ImageClassifierWriter = image_classifier.MetadataWriter\n",
        "# Create the metadata writer.\n",
        "writer = ImageClassifierWriter.create_from_metadata_info(\n",
        "    model_buffer, general_md, input_md, output_md)\n",
        "\n",
        "# Verify the metadata generated by the metadata writer.\n",
        "print(writer.get_metadata_json())\n",
        "\n",
        "# Populate the metadata into the model.\n",
        "writer_utils.save_file(writer.populate(), _SAVE_TO_PATH)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z78vuu6np5sb"
      },
      "source": [
        "## Read the metadata populated to your model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DnWt-4oOselo"
      },
      "source": [
        "You can display the metadata and associated files in a TFLite model through the following code:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5D13YPUsp5VT"
      },
      "outputs": [],
      "source": [
        "from tflite_support import metadata\n",
        "\n",
        "displayer = metadata.MetadataDisplayer.with_model_file(\"mobilenet_v2_1.0_224_metadata.tflite\")\n",
        "print(\"Metadata populated:\")\n",
        "print(displayer.get_metadata_json())\n",
        "\n",
        "print(\"Associated file(s) populated:\")\n",
        "for file_name in displayer.get_packed_associated_file_list():\n",
        "  print(\"file name: \", file_name)\n",
        "  print(\"file content:\")\n",
        "  print(displayer.get_associated_file_buffer(file_name))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "Mq-riZs-TJGt"
      ],
      "name": "Metadata Writer tutorial",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
